package task

import (
	"errors"
	"fmt"
	"proxy_resource_test/services/mill/task/handlers"
	"sync/atomic"
	"time"

	"github.com/ThreeDotsLabs/watermill/message"
	"go.uber.org/zap"
	"proxy_resource_test/utils/logging"
)

const (
	MaxBatchSize   = 300
	MaxWaitTime    = 5
	WorkerPoolSize = 10
	BufferSize     = 10000
)

type ProcessorConfig struct {
	BatchSize int
	WaitTime  int // in seconds
}

var TopicConfigs = map[string]ProcessorConfig{
	"ip_info_region": {BatchSize: MaxBatchSize, WaitTime: MaxWaitTime},
}

type BatchProcessor struct {
	messagesCh chan *message.Message
	workerPool chan struct{}
	pending    int64
	handlers   map[string]func() handlers.BatchHandler
	config     ProcessorConfig
}

func NewBatchProcessor(config ProcessorConfig) *BatchProcessor {
	bp := &BatchProcessor{
		messagesCh: make(chan *message.Message, BufferSize),
		workerPool: make(chan struct{}, WorkerPoolSize),
		handlers:   registerHandlers(),
		config:     config,
	}
	go bp.batchWorker()
	return bp
}

func registerHandlers() map[string]func() handlers.BatchHandler {
	return map[string]func() handlers.BatchHandler{
		"ip_info_region": func() handlers.BatchHandler { return &handlers.IpInfoHandler{} },
	}
}

func (bp *BatchProcessor) AcceptConsumeTask(msg *message.Message) error {
	select {
	case bp.messagesCh <- msg:
		atomic.AddInt64(&bp.pending, 1)
		return nil
	default:
		msg.Nack()
		return errors.New("消息通道已满")
	}
}

func (bp *BatchProcessor) batchWorker() {
	var batch []*message.Message
	timer := time.NewTimer(time.Duration(bp.config.WaitTime) * time.Second)
	defer timer.Stop()

	for {
		select {
		case msg := <-bp.messagesCh:
			batch = append(batch, msg)
			if len(batch) >= bp.config.BatchSize {
				bp.flush(&batch, timer)
			}
		case <-timer.C:
			if len(batch) > 0 {
				bp.flush(&batch, timer)
			}
			timer.Reset(time.Duration(bp.config.WaitTime) * time.Second)
		}
	}
}

func (bp *BatchProcessor) flush(batch *[]*message.Message, timer *time.Timer) {
	toFlush := make([]*message.Message, len(*batch))
	copy(toFlush, *batch)
	*batch = (*batch)[:0]
	timer.Reset(MaxWaitTime)

	go func(msgs []*message.Message) {
		bp.workerPool <- struct{}{}
		defer func() { <-bp.workerPool }()
		if err := bp.processBatch(msgs); err != nil {
			bp.nackMessages(msgs)
		} else {
			bp.ackMessages(msgs)
		}
		atomic.AddInt64(&bp.pending, -int64(len(msgs)))
	}(toFlush)
}

func (bp *BatchProcessor) processBatch(msgs []*message.Message) error {
	if len(msgs) == 0 {
		return nil
	}
	topic := msgs[0].Metadata.Get("topic")
	factory, ok := bp.handlers[topic]
	if !ok {
		return fmt.Errorf("未注册的 topic: %s", topic)
	}

	handler := factory()
	for _, msg := range msgs {
		if err := handler.UnmarshalAndCollect(msg); err != nil {
			logging.Warn(topic, "--解码失败--", zap.Error(err))
			continue
		}
	}

	if err := handler.Handle(); err != nil {
		return fmt.Errorf("处理失败: %w", err)
	}
	return nil
}

func (bp *BatchProcessor) ackMessages(msgs []*message.Message) {
	for _, msg := range msgs {
		_ = msg.Ack()
	}
}

func (bp *BatchProcessor) nackMessages(msgs []*message.Message) {
	for _, msg := range msgs {
		_ = msg.Nack()
	}
}
